# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

import numpy as np

def py_cpu_nms(dets, thresh):
    """Pure Python NMS baseline."""
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]

    return keep



def xywh2xyxy(x):
    y = np.zeros(x.shape)
    y[..., 0] = x[..., 0] - x[..., 2] / 2
    y[..., 1] = x[..., 1] - x[..., 3] / 2
    y[..., 2] = x[..., 0] + x[..., 2] / 2
    y[..., 3] = x[..., 1] + x[..., 3] / 2
    return y

def sigmoid(x):
    s = 1.0 / (1 + np.exp(-x))
    return s

def compute_grid_offsets(grid_size, scaled_anchors):
    stride = 416 / grid_size
    # 32 = imdim / grid
    # size
    grid_x = np.repeat([np.array([i for i in range(grid_size)])], grid_size, axis=0).reshape(1, 1, grid_size, grid_size)
    grid_y = np.repeat([np.array([i for i in range(grid_size)])], grid_size, axis=0).transpose(1, 0).reshape(1, 1, grid_size, grid_size)
    #
    #
    anchor_w = scaled_anchors[:, 0:1].reshape((1, 3, 1, 1))
    anchor_h = scaled_anchors[:, 1:2].reshape((1, 3, 1, 1))
    return grid_x, grid_y, anchor_w, anchor_h

def yolo_output(input, grid_size, scaled_anchors):
    prediction = np.resize(input, (1, 3, 85, grid_size, grid_size)).transpose(0, 1, 3, 4, 2)
    stride = 416 / grid_size
    # Get outputs
    x = sigmoid(prediction[..., 0])  # Center x
    y = sigmoid(prediction[..., 1])  # Center y
    w = prediction[..., 2]  # Width
    h = prediction[..., 3]  # Height
    pred_conf = sigmoid(prediction[..., 4])  # Conf
    pred_cls = sigmoid(prediction[..., 5:])  # Cls pred.

    # If grid size does not match current we compute new offsets
    grid_x, grid_y, anchor_w, anchor_h = compute_grid_offsets(grid_size, scaled_anchors)

    # Add offset and scale with anchors
    pred_boxes = np.zeros(prediction[..., :4].shape)
    pred_boxes[..., 0] = x + grid_x
    pred_boxes[..., 1] = y + grid_y
    pred_boxes[..., 2] = np.exp(w) * anchor_w
    pred_boxes[..., 3] = np.exp(h) * anchor_h

    pred_boxes = pred_boxes.reshape(1, -1, 4) * stride
    pred_conf = pred_conf.reshape(1, -1, 1)
    pred_cls = pred_cls.reshape(1, -1, 80)

    output = np.concatenate((pred_boxes, pred_conf, pred_cls), axis=-1)

    return output
    
    